package com.lijing.huffmanCode;


import java.io.*;
import java.util.*;

/**
 * @Description 赫夫曼编码
 * @title: HuffmanCode
 * @Author LiJing
 * @Date: 2021/2/258:44 上午
 * @Version 1.0
 */
public class HuffmanCode {
    private Node root;
    private Map<Byte, String> huffmanCodes = new HashMap<Byte, String>();
    private StringBuilder stringBuilder = new StringBuilder();

    public byte[] huffmanZip(byte[] bytes){
        /*
         * @Date: 2021/3/25 9:10 上午
         * Step 1: 将bytes转换为node链表
         */
        List<Node> nodeList = getNode(bytes);
        /*
         * @Date: 2021/3/25 9:12 上午
         * Step 2: 生成huffmanTree
         */
        Node huffmanRoot = creatHuffmanTree(nodeList);
        this.setRoot(huffmanRoot);
        /*
         * @Date: 2021/3/25 9:14 上午
         * Step 3: 构建Huffman编码
         */
        getCodes();
        /*
         * @Date: 2021/3/25 9:17 上午
         * Step 4: 打包输出
         */
        return zip(bytes,huffmanCodes);
    }

    /**
     * 方法重载
     */
    public Map<Byte, String> getCodes(){
        getCodes(this.getRoot(), "", stringBuilder);
        return huffmanCodes;

    }

    /**
     * 获取传入的node节点的所有叶子节点的huffman编码，并放入到huffmanCodes集合
     * @param node 传入节点
     * @param code 路径：左子节点为0，右子节点为1
     * @param stringBuilder 用于拼接路径
     */
    public void getCodes(Node node, String code, StringBuilder stringBuilder){
        StringBuilder stringBuilder1 = new StringBuilder(stringBuilder);
        stringBuilder1.append(code);
        if (node != null){
            if (node.getData() == null){
                getCodes(node.getLeft(),"0",stringBuilder1);
                getCodes(node.getRight(),"1",stringBuilder1);
            }else {
                huffmanCodes.put(node.getData(), stringBuilder1.toString());
            }
        }
    }

    /**
     *
     * @param bytes 原始待转化的byte数组
     * @param huffmanCodes 生成的赫夫曼编码map
     * @return 返回
     */
    public byte[] zip(byte[] bytes, Map<Byte, String> huffmanCodes){

        /*
         * @Date: 2021/3/24 10:18 下午
         * Step 1: 将bytes数组按照huffmanCodes转换为100101010101形式,刚在stringBuilder中
         */
        StringBuilder stringBuilder = new StringBuilder();
        for (byte b : bytes){
            stringBuilder.append(huffmanCodes.get(b));
        }
        /*
         * @Date: 2021/3/24 10:19 下午
         * Step 2: 将stringBuilder转换成每8位一个byte
         * 先要计算转换后的huffmanCodesBytes数组的长度，然后每隔8位进行一次转换放入到huffmanCodesBytes中
         */
        int len = (stringBuilder.length() + 7)/8;
        byte[] huffmanCodesBytes = new byte[len];
        int index = 0;
        for (int i = 0; i < stringBuilder.length(); i += 8) {
            byte byteNode;
            if (stringBuilder.length() < i+8){
                byteNode = (byte) Integer.parseInt(stringBuilder.substring(i),2);
            }else {
                byteNode = (byte) Integer.parseInt(stringBuilder.substring(i,i+8),2);
            }

            huffmanCodesBytes[index] = byteNode;
            index++;
        }
        return huffmanCodesBytes;
    }

    /**
     * 将src路径文件压缩到dst
     * @param src 待压缩的文件原路径
     * @param dst 压缩后的文件输出路径
     */
    public void zipFile(String src, String dst){
        FileInputStream is = null;
        OutputStream os = null;
        ObjectOutputStream oos = null;
        try {
            /*
             * @Date: 2021/3/26 1:15 下午
             * Step 1: 创建输入流，读取文件
             */
            is = new FileInputStream(src);
            byte[] bytes = new byte[is.available()];
            is.read(bytes);
            /*
             * @Date: 2021/3/26 1:16 下午
             * Step 2: 进行huffman压缩
             */
            byte[] huffmanBytes = huffmanZip(bytes);
            /*
             * @Date: 2021/3/26 1:19 下午
             * Step 3: 创建输出流，将处理后的huffmanBytes输出到目的路径
             */
            os = new FileOutputStream(dst);
            oos = new ObjectOutputStream(os);
            oos.writeObject(huffmanBytes);
            /*
             * @Date: 2021/3/26 1:24 下午
             * Step 4: 将huffmanCodes也写入到压缩文件
             */
            oos.writeObject(this.getHuffmanCodes());
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }finally {
            if (is != null) {
                try {
                    is.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (os != null) {
                try {
                    os.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (oos != null) {
                try {
                    oos.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     *
     * @param zipFile file to be unzip
     * @param dstFile dst
     */
    public void unZipFile(String zipFile, String dstFile){
        InputStream is = null;
        ObjectInputStream ois = null;
        OutputStream os = null;

        try {
            /*
             * @Date: 2021/3/26 1:39 下午
             * Step 1: 创建文件输入流
             */
            is = new FileInputStream(zipFile);
            ois = new ObjectInputStream(is);
            /*
             * @Date: 2021/3/26 1:44 下午
             * Step 2: 读取huffman编码后的结果和Huffman编码表
             */
            byte[] huffmanBytes = (byte[])ois.readObject();
            Map<Byte, String> huffmanCodes = (Map<Byte, String>)ois.readObject();
            /*
             * @Date: 2021/3/26 1:44 下午
             * Step 3: 解码
             */
            byte[] bytes = decode(huffmanCodes,huffmanBytes);
            /*
             * @Date: 2021/3/26 1:44 下午
             * Step 4: 写数据
             */
            os = new FileOutputStream(dstFile);
            os.write(bytes);

        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        } finally {
            if (is != null) {
                try {
                    is.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 根据byte数组返回对应的节点列表，每个节点由数据和权重决定
     * @param bytes 数组
     * @return 返回list
     */
    public static List<Node> getNodes(byte[] bytes){
        List<Node> listNodes = new ArrayList<Node>();
        Arrays.sort(bytes);
        byte first = 0;
        byte end = 0;
        while (end < bytes.length) {
            while (end < bytes.length && bytes[first] == bytes[end]) {
                end++;
            }
            listNodes.add(new Node(bytes[first], end - first));
            first = end;
        }

        return listNodes;
    }


    /**
     *
     * @param bytes 需要编码的byte数组
     * @return 返回节点链表
     */
    public static List<Node> getNode(byte[] bytes){
        /*
         * @Date: 2021/3/5 7:36 下午
         * Step 1: 创建一个ArrayList
         */
        ArrayList<Node> nodes = new ArrayList<Node>();
        /*
         * @Date: 2021/3/5 7:37 下午
         * Step 2: 遍历bytes，统计每一个byte出现的次数，map[key, value]
         */
        Map<Byte, Integer> counts = new HashMap<>();
        for (byte b : bytes){
            counts.merge(b, 1, Integer::sum);
        }
        /*
         * @Date: 2021/3/6 2:15 下午
         * Step 3: 将map中的元素转换为Node
         */
        for (Map.Entry<Byte, Integer> entry : counts.entrySet()){
            nodes.add(new Node(entry.getKey(), entry.getValue()));
        }
        return nodes;
    }

    /**
     * 生成赫夫曼树
     * @param nodes 待生成赫夫曼树的节点list
     * @return 返回huffmanTree的跟节点
     */
    public static Node creatHuffmanTree(List<Node> nodes){
        while (nodes.size() > 1){
            Collections.sort(nodes);
            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);
            Node parent = new Node(null,leftNode.getWeight()+ rightNode.getWeight());
            parent.setLeft(leftNode);
            parent.setRight(rightNode);
            nodes.remove(leftNode);
            nodes.remove(rightNode);
            nodes.add(parent);
        }
        return nodes.get(0);
    }

    /**
     * 将一个byte转成一个二进制字符串，按照补码的方式保存
     * @param flag 标志是否需要补高位(针对的是最后一个字节)，如果是true，表示需要补高位，如果是false表示不补
     * @param b 待转换的byte
     * @return 返回b对应的补码 如-88  --->  10101000
     */
    public static String byteToBitString(boolean flag, byte b){
        int temp = b;
        if (flag){
            temp |= 256;
        }
        String str = Integer.toBinaryString(temp);
        if (flag){
            return str.substring(str.length()-8);
        }else {
            return str;
        }

    }

    public static byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes){
        /*
         * @Date: 2021/3/25 10:30 上午
         * Step 1: 得到huffmanBytes对应二进制的字符串
         */
        StringBuilder stringBuilder = new StringBuilder();
        for (int i = 0; i < huffmanBytes.length; i++) {
            boolean flag = (i == huffmanBytes.length-1);
            stringBuilder.append(byteToBitString(!flag, huffmanBytes[i]));
        }
        System.out.println(stringBuilder.toString());
        /*
         * @Date: 2021/3/25 10:43 上午
         * Step 2: 把字符串按照huffmanCodes进行解码
         */
        Map<String, Byte> map = new HashMap<>();
        for (Map.Entry<Byte, String> set : huffmanCodes.entrySet()){
            map.put(set.getValue(), set.getKey());
        }
        System.out.println(map);

        List<Byte> list = new ArrayList<>();
        for (int i = 0; i < stringBuilder.length();) {
             int count = 1;
             boolean flag = true;
             Byte b = null;
             while (flag){
                 String key = stringBuilder.substring(i, count+i);
                 b = map.get(key);
                 if (b == null){
                     count++;
                 }else {
                     flag = false;
                 }

             }
            list.add(b);
             i += count;
        }
        byte[] bytes = new byte[list.size()];
        for (int i = 0; i < bytes.length; i++) {
            bytes[i] = list.get(i);
        }
        return bytes;

    }


    public void preOrder(){
        if (root != null){
            root.preOrder();
        }else {
            System.out.println("The huffmanTree is empty!");
        }
    }

    public HuffmanCode(Node root) {
        this.root = root;
    }

    public HuffmanCode() {
    }

    public Node getRoot() {
        return root;
    }

    public void setRoot(Node root) {
        this.root = root;
    }

    public Map<Byte, String> getHuffmanCodes() {
        return huffmanCodes;
    }

    public void setHuffmanCodes(Map<Byte, String> huffmanCodes) {
        this.huffmanCodes = huffmanCodes;
    }

    public StringBuilder getStringBuilder() {
        return stringBuilder;
    }

    public void setStringBuilder(StringBuilder stringBuilder) {
        this.stringBuilder = stringBuilder;
    }
}
